from keras.models import Sequential
from keras.preprocessing.image import ImageDataGenerator
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from keras.callbacks import ModelCheckpoint, TensorBoard
from keras.optimizers import Adam, RMSprop
import cv2
import glob

def cnn(n_classes=1, input_shape=(255, 255, 3)):
    # 创建一个顺序模型
    model = Sequential()

    # 添加卷积层，32 个滤波器，每个滤波器大小为 3x3，使用 ReLU 激活函数，输入形状为图像的尺寸和通道数
    model.add(Conv2D(32, (3, 3), activation='relu', input_shape=input_shape))

    # 添加最大池化层，池化窗口大小为 2x2
    model.add(MaxPooling2D((2, 2)))

    # 再添加一个卷积层和最大池化层
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D((2, 2)))

    # 再添加一个卷积层和最大池化层
    model.add(Conv2D(128, (3, 3), activation='relu'))
    model.add(MaxPooling2D((2, 2)))

    model.add(Conv2D(32, (3, 3), activation='relu'))
    # 将多维的输出展平为一维
    model.add(Flatten())

    # 添加全连接层，32 个神经元，使用 ReLU 激活函数
    model.add(Dense(32, activation='relu'))

    # 添加输出层，假设是分类问题，有 1 个类别，使用 sigmoid 激活函数
    model.add(Dense(1, activation='sigmoid' if n_classes==1 else 'softmax'))
    return model

image_path = 'deep_net\\test'
img = cv2.imread((glob.glob(image_path+"\\image\\*.jpg")+glob.glob(image_path+"\\image\\*.png"))[0])
datagen = ImageDataGenerator(rescale=1./255,rotation_range=10,width_shift_range=0.2,height_shift_range=0.2,
    shear_range=0.2,zoom_range=0.2,horizontal_flip=True,fill_mode='constant')
generate_data = datagen.flow_from_directory(image_path,target_size=img.shape[:2],batch_size=1)       #样本数不易过大

model = cnn(1,img.shape)
model.summary()
# 编译模型，指定损失函数、优化器和评估指标
model.compile(optimizer=Adam(learning_rate=0.001), loss='binary_crossentropy', metrics=['accuracy'])
model.fit(generate_data, steps_per_epoch=2, epochs=10, callbacks=
    [
        TensorBoard('deep_net/logs',write_images=True,histogram_freq=1,write_graph=False),
        ModelCheckpoint("deep_net\\param\\cnn.keras"),
    ])
